% Computes the action for a given state.
% Computes: aileron, elevator, rudder, and collective controls.
function action = policy(x, dx, w, use_linear)
    % Result of LQR copy pasta-ed here.
    K_ss = [-0.279551627722307,0,0,-0.000570795136247815,0,0,0,0,0,-0.0491011200797148,-0.00281723788346760,0,-0.0267969669504523,-0.00153978392447130,-0.0719638984766877,0,0,-0.725150000875433,0,0,-2.77227796006990e-18;0,-0.262032755060953,-0.0128687138715770,0,0,0,0,0,-0.0540325889417506,0,0,-0.0270381364143380,0,0,0,0.0899168482557423,-0.00467829008114874,0,0.829495035324250,-0.0456460689451671,0;0,-0.0128687138715770,-0.139011060463097,0,0,0,0,0,-0.00463303057597533,0,0,-0.00236675497001590,0,0,0,0.00651234341647562,-0.0135693045830338,0,0.0659108710191011,-0.0331326669194076,0;-0.000570795136247816,0,0,-0.289174206017467,0,0,0,0,0,-0.00204577520853476,0.0362689557401554,0,-0.00150823526778955,0.0266186216840679,-0.000628442544581799,0,0,-0.0114338201872748,0,0,2.98925714219746e-17];

    % Compute the action that we are going to take.
    if use_linear
        % This is a linear function of our features.
        action = [K_ss(:, 1:8) w' K_ss(:, 21)]*dx;
    else
        % This is a nonlinear function based on some neural net.
        euler_angles = q_to_euler(x(18:21));
        euler_angles_delta = q_to_euler(dx(18:21));
        
        xyz = express_vector_in_quat_frame(x(12:14), x(18:21));
        uvw = express_vector_in_quat_frame(x(9:11), x(18:21));

        % Put into easy to index matrix.
        f = [xyz(1);  % n - n_desired,  errx
             xyz(2);  % e - e_desired,  erry
             xyz(3);  % d - d_desired,  errz
             uvw(1);    % n_dot,  x_dot
             uvw(2);   % e_dot,  y_dot
             uvw(3);   % d_dot,  z_dot
             euler_angles(2);  % pitch,  theta
             euler_angles(1);  % roll,  phi
             x(17);   % yaw_dot,  omega_dot
             euler_angles_delta(3); % yaw - yaw_desired, err_omega
             1];

         t1 = w(1)*f(11) + w(2)*tanh(w(3)*f(1)) + w(4)*f(1) + w(5)*f(4) + w(6)*f(7);
         t2 = w(10)*f(11) + w(11)*tanh(w(12)*f(2)) + w(13)*f(2) + w(14)*f(5) + w(15)*f(8);
         t3 = w(19)*f(11) + w(20)*tanh(w(21)*f(3)) + w(22)*f(3) + w(23)*f(6);
         t4 = w(27)*f(9) + w(28)*f(10);

         a1 = w(7)*tanh(w(8)*t1) + w(9)*t1;
         a2 = w(16)*tanh(w(17)*t2) + w(18)*t2;
         a3 = w(24)*tanh(w(25)*t3) + w(26)*t3;
         a4 = w(29)*tanh(w(30)*t4) + w(31)*t4;

         action = [a2; a1; a4; a3];
    end
end